# -*- coding: utf-8 -*-
import base64
import datetime
import hashlib
import hmac
import json
import random
import requests
import os


class LangboatOpenClient:
    """澜舟开放平台客户端"""

    def __init__(self,
                 access_key: str,
                 access_secret: str,
                 url: str = "https://open.langboat.com"):
        self.access_key = access_key
        self.access_secret = access_secret
        self.url = url

    def _build_header(self, query: dict, data: dict) -> dict:
        accept = "application/json"
        # 1. body MD5 加密
        content_md5 = base64.b64encode(
            hashlib.md5(
                json.dumps(data).encode("utf-8")
            ).digest()
        ).decode()
        content_type = "application/json"
        gmt_format = '%a, %d %b %Y %H:%M:%S GMT'
        date = datetime.datetime.utcnow().strftime(gmt_format)
        signature_method = "HMAC-SHA256"
        signature_nonce = str(random.randint(0, 65535))
        header_string = f"POST\n{accept}\n{content_md5}\n{content_type}\n" \
                        f"{date}\n{signature_method}\n{signature_nonce}\n"

        # 2. 计算 queryToSign
        queries_str = []
        for k, v in sorted(query.items(), key=lambda item: item[0]):
            if isinstance(v, list):
                for i in v:
                    queries_str.append(f"{k}={i}")
            else:
                queries_str.append(f"{k}={v}")
        queries_string = '&'.join(queries_str)
        # 3.计算 stringToSign
        sign_string = header_string + queries_string
        # 4.计算 HMAC-SHA256 + Base64
        secret_bytes = self.access_secret.encode("utf-8")
        # 5.计算签名
        signature = base64.b64encode(
            hmac.new(secret_bytes, sign_string.encode("utf-8"), hashlib.sha256).digest()
        ).decode()
        res = {
            "Content-Type": content_type,
            "Content-MD5": content_md5,
            "Date": date,
            "Accept": accept,
            "X-Langboat-Signature-Method": signature_method,
            "X-Langboat-Signature-Nonce": signature_nonce,
            "Authorization": f"{self.access_key}:{signature}"
        }
        return res

    def inference(self, queries: dict, data: dict) -> (int, dict):
        """
        服务调用: POST https://open.langboat.com/chat
        :param queries: query 参数
        :param data: request body 数据
        :return: response status, response body to json
        """
        headers = self._build_header(queries, data)
        response = requests.post(url=self.url + "/chat",
                                 headers=headers, params=queries, json=data)
        return response.status_code, response.json()


if __name__ == '__main__':
    _access_key = os.getenv("MCHAT_ACCESS_KEY", "")
    _access_secret = os.getenv("MCHAT_ACCESS_SECRET", "")

    client = LangboatOpenClient(
        access_key=_access_key,
        access_secret=_access_secret
    )

    # 孟子GPT
    _queries = {}
    _data = {
        "model": "mengzi-lite",
        "messages": [
            {
                "role": "user",
                "content": "你知道牛顿吗"
            }
        ],
        "n": 1,
        "max_tokens": 1024,
        "user": ""
    }
    status_code, result = client.inference(_queries, _data)
    print("response status:", status_code)
    print("response json:", json.dumps(result, ensure_ascii=False, indent=2))
